package entity

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/grafov/m3u8"
	"strconv"
	"strings"
)

var PngSign = []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A}

// TsInfo 用于保存 ts 文件的下载地址和文件名
type TsInfo struct {
	*m3u8.MediaSegment
	*EncryptInfo
}

type EncryptInfo struct {
	Method string
	Key    []byte
	Iv     []byte
	block  cipher.Block
}

func NewEncryptInfo(method, ivs string, key []byte) (*EncryptInfo, error) {
	var iv []byte
	var err error
	if ivs != "" {
		iv, err = hex.DecodeString(strings.TrimPrefix(ivs, "0x"))
		if err != nil {
			return nil, err
		}
	}
	encInfo := &EncryptInfo{
		Method: method,
		Key:    key,
		Iv:     iv,
	}
	block, err := aes.NewCipher(encInfo.Key)
	encInfo.block = block
	return encInfo, err
}

func (e *EncryptInfo) Decode(seq uint64, encrypted []byte) ([]byte, error) {
	iv := e.Iv
	if len(iv) == 0 {
		if e.Method == EncryptMethod_AES128 {
			iv = make([]byte, 16)
			binary.BigEndian.PutUint64(iv[8:], seq)
		} else {
			return nil, errors.New("AesDecrypt method not support " + strconv.Quote(e.Method))
		}
	}
	blockMode := cipher.NewCBCDecrypter(e.block, iv)
	if len(iv) == 0 || len(encrypted)%len(iv) != 0 {
		return nil, errors.New("AesDecrypt invalid encrypted data len " + strconv.Itoa(len(encrypted)))
	}
	origData := make([]byte, len(encrypted))
	blockMode.CryptBlocks(origData, encrypted)
	length := len(origData)
	unpadding := int(origData[length-1])
	if length-unpadding < 0 {
		return nil, fmt.Errorf(`invalid length of unpadding %v - %v`, length, unpadding)
	}
	return origData[:(length - unpadding)], nil

}

// https://datatracker.ietf.org/doc/html/rfc8216#section-4.3.2.4
const (
	EncryptMethod_NONE   = `NONE`
	EncryptMethod_AES128 = `AES-128`
)

func M3u8Parse(m3u8Bytes []byte) (m3u8.Playlist, m3u8.ListType, error) {
	return m3u8.Decode(*bytes.NewBuffer(m3u8Bytes), true)
}
